##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)
library(ggpubr)
library(patchwork)
library(ggsci)
library(rstatix)
library(scales)

##########################################################################################

option_list <- list(
    make_option(c("--info_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic/"
    info_file <- paste(work_dir,"/baseTable/STAD_Info.addBurden.MSI_MSS.addCNVType.tsv",sep="")
    images_path <- paste(work_dir,"/images/mutBurden",sep="")

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

info_file <- opt$info_file
images_path <- opt$images_path

dir.create(images_path , recursive = T)

###########################################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")

col_im <- brewer.pal(9,"YlGnBu")[6:8]
names(col_im) <- c("IM(IGC)" , "IM(DGC)" , "IM(IGC_DGC)")

Variant_Types <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################

dat_sample <- data.frame(fread( info_file ))

###########################################################################################
## 多个样本负荷取中位数
burden_type <- "All"
dat_plot2 <- dat_sample %>%
    group_by( Patient , Class , Type , TCGA_Class , Age ) %>%
    summarize( BurdenExon = median(BurdenExon) )
dat_plot2$Age_divide <- ifelse(dat_plot2$Age > 66 , "Older" , "Younger")

###########################################################################################

msi_sample <- unique(subset( dat_sample , TCGA_Class == "MSI" )$Patient)
gs_sample <- unique(subset( dat_sample , TCGA_Class == "GS" )$Patient)
cin_sample <- unique(subset( dat_sample , TCGA_Class == "CIN" )$Patient)

###########################################################################################

trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("p == ",down," %*% 10","^",up)
    return(text)
}

###########################################################################################

plotBurden <- function(dat_plot_tmp = dat_plot_tmp , type = type ){

    dat_plot_tmp$title <- type
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% msi_sample , "MSI/POLE" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% gs_sample , "GS" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% cin_sample , "CIN" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- factor( dat_plot_tmp$TCGA_Class , levels = c("MSI/POLE" , "CIN" , "GS") , order = T )

    res_tmp <- dat_plot_tmp %>%
    group_by( TCGA_Class ) %>%
    summarize( BurdenExon = median(BurdenExon) )
    out_name <- paste0( images_path , "/mutBurden.IM.TCGA_Type.GS_CIN_MSI.",burden_type,".All.tsv" )  
    write.table( res_tmp , out_name , row.names = F , sep = "\t" , quote = F )

    sample_num <- dat_plot_tmp %>%
    group_by( TCGA_Class ) %>%
    summarize( nums = length(unique(Patient)) )

    dat_plot_tmp <- merge( dat_plot_tmp , sample_num , by = "TCGA_Class" )
    dat_plot_tmp$TCGA_Class_num <- paste0(dat_plot_tmp$TCGA_Class , "\n" , "(" , dat_plot_tmp$nums , ")" )

    my_comparisons_1 <- list( c(1,2) , c(1,3) , c(2,3))

    #stat.test <- dat_plot_tmp %>% 
    #  wilcox_test(BurdenExon ~ TCGA_Class_num , comparisons = my_comparisons_1) %>%
    #  add_xy_position(x = "TCGA_Class_num") %>%
    #  mutate(myformatted.p = paste0("p = ", p))

    #print(stat.test)

    dat_plot_tmp$TCGA_Class <- factor( dat_plot_tmp$TCGA_Class , levels = c( "GS" , "CIN" , "MSI/POLE" ) , order = T )

    col_use <- c(rgb(red=179,green=34,blue=35,alpha=255,max=255) ,
        rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
        rgb(red=2,green=100,blue=190,alpha=255,max=255) 
    )

    names(col_use) <- c("GS" , "MSI/POLE" , "CIN")

    plot <- ggplot( dat_plot_tmp , aes( x = TCGA_Class , y = BurdenExon , color = TCGA_Class ) ) +
        geom_line( aes( group = Patient ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
        geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
        geom_jitter(position = position_jitter(0.2) , size = 1 , alpha = 1) +
        scale_color_manual(values = col_use) +
        facet_grid(.~title) +
        scale_y_continuous(
                trans = sqrt_trans(),
                breaks = c(1:5)
                ) + 
        xlab(NULL) +
        ylab("Mutation rate per MB")+
        #ylim(0,5.5)+
        theme_bw() +
        stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
        theme(
            legend.position = 'none',
            legend.title = element_blank() ,
            panel.grid.major=element_blank(),
            panel.grid.minor=element_blank(),
            panel.background = element_blank(),
            panel.border = element_blank(),
            plot.title = element_text(size = 12,color="black",face='bold'),
            legend.text = element_text(size = 12,color="black",face='bold'),
            axis.text.y = element_text(size = 12,color="black",face='bold'),
            axis.title.x = element_text(size = 12,color="black",face='bold'),
            axis.title.y = element_text(size = 12,color="black",face='bold'),
            axis.text.x = element_text(size = 12,color="black",face='bold') ,
            axis.ticks.length = unit(0.2, "cm") ,
            strip.text.x = element_text(size = 15, colour = "black",face='bold') ,
            axis.line = element_line(size = 0.5)) 
    return(plot)
}

###########################################################################################
## 分IM+IGC、IM+DGC分别比较
type <- "IM"
dat_plot_tmp <- data.frame(dat_plot2)
dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" & Type != "IM + IGC + DGC")
for(age_d in c("Younger" , "Older")){
    p1 <- plotBurden(dat_plot_tmp = subset(dat_plot_tmp , Age_divide == age_d) , type = type )
    out_name <- paste0( images_path , "/mutBurden.IM.TCGA_Type.GS_CIN_MSI.",burden_type,"." , age_d ,".pdf" )  
    ggsave(file=out_name,plot=p1,width=3.8/1.2,height=4.4/1.2)
}